import os
import sys
import math
import collections
import argparse
from common import BinaryFile


class Vector:
  def __init__(self,x = 0,y = 0,z = 0):
    self.x = x
    self.y = y
    self.z = z

  def __neg__(self): return Vector(-self.x,-self.y,-self.z)
  def __add__(self,v): return Vector(self.x + v.x,self.y + v.y,self.z + v.z)
  def __sub__(self,v): return Vector(self.x - v.x,self.y - v.y,self.z - v.z)
  def __mul__(self,f): return Vector(self.x*f,self.y*f,self.z*f)
  def __rmul__(self,f): return Vector(f*self.x,f*self.y,f*self.z)
  def __div__(self,f): return Vector(self.x/f,self.y/f,self.z/f)

  def norm_sq(self): return self.x*self.x + self.y*self.y + self.z*self.z
  def norm(self): return math.sqrt(self.norm_sq())
  def unit(self): return self/self.norm()

def dot(a,b): return a.x*b.x + a.y*b.y + a.z*b.z
def cross(a,b):
  return Vector(a.y*b.z - a.z*b.y,a.z*b.x - a.x*b.z,a.x*b.y - a.y*b.x)

#------------------------------------------------------------------------------

class Triangle:
  def __init__(self,u,v,w):
    self.u = u
    self.v = v
    self.w = w
    self.n = cross(v - u,w - u).unit()

class Cube:
  def __init__(self,base,width):
    self.hw = width/2.0
    self.c = base + Vector(self.hw,self.hw,self.hw)

def tricube_overlap(t,c):
  '''Intersection test for triangle and axis-aligned cube.
     http://jgt.akpeters.com/papers/AkenineMoller01/tribox.html'''
  
  def axis_test(a1,a2,b1,b2,c1,c2):
    p = a1*b1 + a2*b2
    q = a1*c1 + a2*c2
    r = c.hw*(abs(a1) + abs(a2))
    return min(p,q) > r or max(p,q) < -r
  
  v0 = t.u - c.c
  v1 = t.v - c.c
  v2 = t.w - c.c

  if min(v0.x,v1.x,v2.x) > c.hw or max(v0.x,v1.x,v2.x) < -c.hw: return False
  if min(v0.y,v1.y,v2.y) > c.hw or max(v0.y,v1.y,v2.y) < -c.hw: return False
  if min(v0.z,v1.z,v2.z) > c.hw or max(v0.z,v1.z,v2.z) < -c.hw: return False

  d = dot(t.n,v0)
  r = c.hw*(abs(t.n.x) + abs(t.n.y) + abs(t.n.z))
  if d > r or d < -r: return False

  e = v1 - v0
  if axis_test(e.z,-e.y,v0.y,v0.z,v2.y,v2.z): return False
  if axis_test(-e.z,e.x,v0.x,v0.z,v2.x,v2.z): return False
  if axis_test(e.y,-e.x,v1.x,v1.y,v2.x,v2.y): return False

  e = v2 - v1
  if axis_test(e.z,-e.y,v0.y,v0.z,v2.y,v2.z): return False
  if axis_test(-e.z,e.x,v0.x,v0.z,v2.x,v2.z): return False
  if axis_test(e.y,-e.x,v0.x,v0.y,v1.x,v1.y): return False

  e = v0 - v2
  if axis_test(e.z,-e.y,v0.y,v0.z,v1.y,v1.z): return False
  if axis_test(-e.z,e.x,v0.x,v0.z,v1.x,v1.z): return False
  if axis_test(e.y,-e.x,v1.x,v1.y,v2.x,v2.y): return False

  return True

#------------------------------------------------------------------------------

def read_obj(filename):
  bb_min = Vector(None,None,None)
  bb_max = Vector(None,None,None)
  vertices = []
  triangles = []

  f = file(filename,'r')
  while True:
    line = f.readline()
    if not line: break
    
    elif line[0:1] == 'v' and line[1:2] not in ('t','n','p'):
      v = line[2:].split()
      x = float(v[0])
      y = float(v[1])
      z = float(v[2])
      if x < bb_min.x or bb_min.x is None: bb_min.x = x
      if y < bb_min.y or bb_min.y is None: bb_min.y = y
      if z < bb_min.z or bb_min.z is None: bb_min.z = z
      if x > bb_max.x or bb_max.x is None: bb_max.x = x
      if y > bb_max.y or bb_max.y is None: bb_max.y = y
      if z > bb_max.z or bb_max.z is None: bb_max.z = z
      vertices.append(Vector(x,y,z))
      
    elif line[0:1] == 'f':
      t = line[2:].split()
      if len(t) < 3: continue
      u = vertices[int(t[0].split('/')[0]) - 1]
      v = vertices[int(t[1].split('/')[0]) - 1]
      w = vertices[int(t[2].split('/')[0]) - 1]
      if cross(v - u,w - u).norm_sq() < 0.001: continue # TODO: find a better solution
      triangles.append(Triangle(u,v,w))

  return (triangles,bb_min,bb_max)

#------------------------------------------------------------------------------

class OctreeNode(Cube):
  def __init__(self,base,width,triangles,indices,max_triangles,min_width):
    Cube.__init__(self,base,width)
    self.is_leaf = True
    self.indices = []
    for i in indices:
      if tricube_overlap(triangles[i],self): self.indices.append(i)
    if len(self.indices) > max_triangles and self.hw >= min_width:
      self.is_leaf = False
      self.branches = [OctreeNode(base + self.hw*Vector(i,j,k),self.hw,triangles,self.indices,max_triangles,min_width)
                       for k in xrange(2) for j in xrange(2) for i in xrange(2)]
      del self.indices

  def write(self,f,parent_base):
    pos = f.tell()
    f.seek(0,os.SEEK_END)
    if self.is_leaf:
      f.write_u32_at((f.tell() - parent_base - 2) | 0x80000000,pos)
      for i in self.indices: f.write_u16(i + 1)
      f.write_u16(0)
    else:
      f.write_u32_at(f.tell() - parent_base,pos)
      base = f.tell()
      f.write('\x00'*4*8)
      f.seek(base)
      for node in self.branches: node.write(f,base)
    f.seek(pos + 4)

#------------------------------------------------------------------------------

def next_exponent(value,base):
  '''Return the lowest integer n such that value <= base**n.'''
  if value <= 1: return 0
  return int(math.ceil(math.log(value,base)))

def write_kcl(filename,triangles,bb_min,bb_max,max_triangles,min_width):
  n_min = next_exponent(min_width,2)
  n_x = max(next_exponent(bb_max.x - bb_min.x,2),n_min)
  n_y = max(next_exponent(bb_max.y - bb_min.y,2),n_min)
  n_z = max(next_exponent(bb_max.z - bb_min.z,2),n_min)
  n = max(min(n_x,n_y,n_z),n_min)
  
  divs_x = 2**(n_x - n)
  divs_y = 2**(n_y - n)
  divs_z = 2**(n_z - n)
  width = 2**n

  octree = [OctreeNode(bb_min + width*Vector(i,j,k),width,triangles,xrange(len(triangles)),max_triangles,min_width)
            for k in xrange(divs_z) for j in xrange(divs_y) for i in xrange(divs_x)]

  f = BinaryFile(filename,'w')
  
  # header
  f.write_u32(0) # vertex offset placeholder
  f.write_u32(0) # vector offset placeholder
  f.write_u32(0) # triangle offset placeholder
  f.write_u32(0) # octree offset placeholder
  f.write_u32(122880) # unknown
  f.write_s32(bb_min.x*4096) # x min
  f.write_s32(bb_min.y*4096) # y min
  f.write_s32(bb_min.z*4096) # z min
  f.write_u32((0xFFFFFFFF << n_x) & 0xFFFFFFFF) # x mask
  f.write_u32((0xFFFFFFFF << n_y) & 0xFFFFFFFF) # y mask
  f.write_u32((0xFFFFFFFF << n_z) & 0xFFFFFFFF) # z mask
  f.write_u32(n) # coordinate shift
  f.write_u32(n_x - n) # y shift
  f.write_u32(n_x - n + n_y - n) # z shift
  f.write_u32(102400) # unknown

  # vertex section
  f.write_u32_at(f.tell(),0) # vertex offset
  for t in triangles:
    f.write_s32(t.u.x*4096)
    f.write_s32(t.u.y*4096)
    f.write_s32(t.u.z*4096)

  # vector section
  f.write_u32_at(f.tell(),4) # vector offset
  for t in triangles:
    a = -cross(t.w - t.u,t.n).unit()
    b = cross(t.v - t.u,t.n).unit()
    c = cross(t.w - t.v,t.n).unit()
    f.write_s16(t.n.x*4096)
    f.write_s16(t.n.y*4096)
    f.write_s16(t.n.z*4096)
    f.write_s16(a.x*4096)
    f.write_s16(a.y*4096)
    f.write_s16(a.z*4096)
    f.write_s16(b.x*4096)
    f.write_s16(b.y*4096)
    f.write_s16(b.z*4096)
    f.write_s16(c.x*4096)
    f.write_s16(c.y*4096)
    f.write_s16(c.z*4096)

  while f.tell() % 4 != 0: f.write_u8(0)

  # triangle section
  f.write_u32_at(f.tell() - 0x10,8) # triangle offset
  for i,t in enumerate(triangles):
    c = cross(t.w - t.v,t.n).unit()
    f.write_s32(dot(t.w - t.u,c)*4096) # length
    f.write_u16(i) # vertex index
    f.write_u16(4*i) # normal index
    f.write_u16(4*i + 1) # a index
    f.write_u16(4*i + 2) # b index
    f.write_u16(4*i + 3) # c index
    f.write_u16(0) # surface type

  # octree section
  f.write_u32_at(f.tell(),12) # octree offset
  base = f.tell()
  nodes = collections.deque(octree)
  node_starts = collections.deque([base for node in octree])
  free_address = base + len(octree)*4

  while len(nodes) > 0:
    current = nodes.popleft()
    current_start = node_starts.popleft()
    if current.is_leaf: f.write_u32(0)
    else:
      f.write_u32(free_address - current_start)
      for node in current.branches:
        nodes.append(node)
        node_starts.append(free_address)
      free_address += 0x20

  f.seek(base)
  nodes = collections.deque(octree)
  node_starts = collections.deque([base for node in octree])
  free_address = base + len(octree)*4

  while len(nodes) > 0:
    current = nodes.popleft()
    current_start = node_starts.popleft()
    pos = f.tell()
    if current.is_leaf:
      f.seek(0,os.SEEK_END)
      f.write_u32_at(0x80000000 | (f.tell() - current_start - 2),pos)
      for i in current.indices: f.write_u16(i + 1)
      f.write_u16(0)
    else:
      for node in current.branches:
        nodes.append(node)
        node_starts.append(free_address)
      free_address += 0x20
    f.seek(pos + 4)

#------------------------------------------------------------------------------

parser = argparse.ArgumentParser(description='Convert .obj files to .kcl files.')
parser.add_argument('infile',help='file to convert')
parser.add_argument('outfile',nargs='?',help='output file')
args = parser.parse_args()

in_root,in_ext = os.path.splitext(args.infile)
if args.outfile is None: args.outfile = in_root + '.kcl'

triangles,bb_min,bb_max = read_obj(args.infile)
write_kcl(args.outfile,triangles,bb_min,bb_max,10,1)

#______________________________________________________________________________
